/**
 * 
 */
package edu.arizona.bio5.onto.decomposition.persistent.xml;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.io.Writer;
import java.net.URI;
import java.util.Set;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.semanticweb.owlapi.apibinding.OWLManager;
import org.semanticweb.owlapi.io.OWLXMLOntologyFormat;
import org.semanticweb.owlapi.io.StringDocumentTarget;
import org.semanticweb.owlapi.model.IRI;
import org.semanticweb.owlapi.model.OWLAxiom;
import org.semanticweb.owlapi.model.OWLOntology;
import org.semanticweb.owlapi.model.OWLOntologyManager;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.xml.sax.InputSource;

import uk.ac.manchester.cs.graph.io.GraphMLWriterImpl;
import edu.arizona.bio5.onto.decomposition.Atom;
import edu.arizona.bio5.onto.decomposition.AtomicDecomposition;
import edu.arizona.bio5.onto.decomposition.persistent.ADPersistenceException;
import edu.arizona.bio5.onto.decomposition.persistent.ADPersister;
import edu.arizona.bio5.onto.util.IOUtils;

/**
 * @author Pavel Klinov
 *
 */
public class XMLADPersister implements ADPersister {

	private File m_rootDir = null;
	
	public void setRootDir(final File rootDir) {
		m_rootDir = rootDir;
	}
	
	/* 
	 */
	@Override
	public void persist(AtomicDecomposition ad) throws ADPersistenceException {
		//First, some cleanup
		System.out.println("Preparation");
		prepare();
		//Second, save all the crap
		persistDocument(getADInfoDocument(ad), XMLADConstants.AD_INFO_DOC);
		persistDocument(getGraphDocument(ad), XMLADConstants.GRAPH_DOC);
		//TODO adapt to multiple labels
		//System.out.print("Storing labels...");
		//persistDocument(getLabelsDocument(ad), XMLADConstants.LABELS_DOC);
		//System.out.println("done!");
		System.out.print("Storing term-based index...");
		persistDocument(getTermIndexDocument(ad), XMLADConstants.TERM_INDEX_DOC);
		System.out.println("done!");
		//Finally, serialize axioms
		System.out.print("Storing atoms...");
		persistAtoms(ad);
		persistDocument(getOntologyDocument(ad.getGlobalAxioms(), IRI.create(ad.getOntologyURI() + "/global_axioms")), XMLADConstants.GLOBAL_AXIOMS_DOC);
		System.out.println("done!");
		//Let subclasses do some clean-up, close connections, etc.
		finish();
	}

	protected Document getTermIndexDocument(AtomicDecomposition ad) throws ADPersistenceException {
		Document doc = createDocument();
		Element rootElem = doc.createElement(XMLADConstants.TERM_INDEX_ROOT);
		
		doc.appendChild(rootElem);
		
		if (ad.getTermBasedIndex() == null) return doc;
		
		for (URI term : ad.getTermBasedIndex().keySet()) {
			Element termElem = doc.createElement(XMLADConstants.TERM_ELEM);
			
			termElem.setAttribute(XMLADConstants.TERM_URI_ATTRIBUTE, term.toString());
			
			for (Atom atom : ad.getTermBasedIndex().get(term)) {
				Element atomElem = doc.createElement(XMLADConstants.ATOM_ELEM);
				
				atomElem.setAttribute(XMLADConstants.ATOM_ID_ATTR, atom.getID().toString());
				termElem.appendChild(atomElem);
			}
			
			rootElem.appendChild(termElem);
		}
		
		return doc;
	}

	protected void finish() {}

	protected void prepare() throws ADPersistenceException {
		IOUtils.cleanRecursively(m_rootDir);
		
		if (!m_rootDir.exists()) m_rootDir.mkdirs();

		new File(m_rootDir.getPath() + "/atoms").mkdirs();
	}

	protected void persistDocument(Document doc, String name) throws ADPersistenceException { 
		try {
			File mdFile = IOUtils.createNewFile(m_rootDir.getPath() + "/" + name);
			Writer writer = new FileWriter(mdFile);
			
			writeOutDocument(doc, writer);
			writer.flush();
			writer.close();
			
		} catch (IOException e) {
			throw new ADPersistenceException(e);
		}
	}

	protected void persistAtoms(AtomicDecomposition ad) throws ADPersistenceException {
		for (Atom atom : ad.getAtoms()) {
			persistDocument(getAtomDocument(ad.getOntologyURI().toString(), atom), "atoms/" + XMLADConstants.ATOM_DOC_PREFIX + atom.getID() + ".owl");
		}
	}	
	
	protected Document getAtomDocument(String ontoURI, Atom atom) throws ADPersistenceException {				
		return getOntologyDocument(atom.getAllAxioms(), IRI.create(ontoURI + "/" + XMLADConstants.ATOM_DOC_PREFIX + atom.getID() + ".owl"));
	}
	
	protected Document getOntologyDocument(Set<OWLAxiom> axioms, IRI iri) throws ADPersistenceException {
		try {
			OWLOntologyManager mgr = OWLManager.createOWLOntologyManager();
			StringDocumentTarget target = new StringDocumentTarget();
			OWLOntology ontology = mgr.createOntology(axioms, iri);
			DocumentBuilder builder = DocumentBuilderFactory.newInstance().newDocumentBuilder();
			
			mgr.saveOntology(ontology, new OWLXMLOntologyFormat(), target);
			//Now parse the string back to document
			//TODO Obviously inefficient, writer a better target impl, smth like XMLDocumentTarget
			Document doc = builder.parse(new InputSource(new StringReader(target.toString())));
			
			return doc;
		} catch (Throwable e) {
			throw new ADPersistenceException(e);
		}
	}
	
/*	protected Document getLabelsDocument(AtomicDecomposition ad) throws ADPersistenceException {
		Document doc = createDocument();
		Element rootElem = doc.createElement(XMLADConstants.LABELS_ROOT_ELEM);
		
		doc.appendChild(rootElem);
		
		for (Atom atom : ad.getAtoms()) {
			
			Element atomElem = doc.createElement(XMLADConstants.ATOM_ELEM);
			
			atomElem.setAttribute(XMLADConstants.ATOM_ID_ATTR, atom.getID().toString());
			saveLabel((SeedSigBasedLabel)atom.getLabel(), atomElem, doc);
			rootElem.appendChild(atomElem);
		}
		
		return doc;
	}*/
	
	protected Document getGraphDocument(AtomicDecomposition ad) throws ADPersistenceException {
		Writer writer = new StringWriter();
		GraphMLWriterImpl<Atom, Object> graphWriter = new GraphMLWriterImpl<Atom, Object>();
		
		try {
			graphWriter.write(ad.getGraph(), writer);
		} catch (IOException e) {
			throw new ADPersistenceException(e);
		}
		
		return graphWriter.getLastDocument();
	}
	
	
	protected Document getADInfoDocument(AtomicDecomposition ad) throws ADPersistenceException {
		try {
			Document doc = createDocument();
			Element rootElem = doc.createElement(XMLADConstants.AD_INFO_ROOT_ELEM);
			
			rootElem.setAttribute(XMLADConstants.AD_MODULE_TYPE_ATTR, ad.getModuleType().name());
			
			Element atomNumber = doc.createElement(XMLADConstants.AD_ATOM_NUM_ELEMENT);
			
			atomNumber.appendChild(doc.createTextNode(String.valueOf(ad.getAtoms().size())));
			
			Element ontoUriElem = doc.createElement(XMLADConstants.AD_ONTO_URI_ELEMENT);
			
			ontoUriElem.appendChild(doc.createTextNode(ad.getOntologyURI().toString()));
			
			rootElem.appendChild(atomNumber);
			rootElem.appendChild(ontoUriElem);
			doc.appendChild(rootElem);
			
			return doc;
		} catch (Throwable e) {
			throw new ADPersistenceException(e);
		}	
	}

	private void writeOutDocument(Document doc, Writer writer) throws ADPersistenceException {
		try {
			TransformerFactory transfac = TransformerFactory.newInstance();
			Transformer trans = transfac.newTransformer();
			
			trans.setOutputProperty(OutputKeys.INDENT, "yes");

			StreamResult result = new StreamResult(writer);
			DOMSource source = new DOMSource(doc);
			
			trans.transform(source, result);
		} catch (Exception e) {
			e.printStackTrace();
			throw new ADPersistenceException(e.getMessage());
		}
	}

/*	private void saveLabel(SeedSigBasedLabel label, Element atomElem, Document doc) {
		if (label != null) {
			Element labelElem = doc.createElement(XMLADConstants.LABEL_ELEM);
			
			labelElem.setAttribute(XMLADConstants.COMPLETE_LABEL_ATTRIBUTE, String.valueOf(label.isComplete()));
			
			for (Set<OWLEntity> mss : label.getSeedSignatures()) {
				Element sigElem = doc.createElement(XMLADConstants.SIGNATURE_ELEM);
				
				for (OWLEntity term : mss) {
					Element termElem = doc.createElement(XMLADConstants.TERM_ELEM);
					
					termElem.appendChild(doc.createTextNode(OntologyUtils.termToString(term)));
					
					sigElem.appendChild(termElem);
				}
				
				labelElem.appendChild(sigElem);
			}
			
			atomElem.appendChild(labelElem);
		}
	}*/

	protected Document createDocument() throws ADPersistenceException {
		try {
			return DocumentBuilderFactory.newInstance().newDocumentBuilder().newDocument();
		} catch (ParserConfigurationException e) {
			throw new ADPersistenceException(e.getMessage());
		}
	}
}